'''
Created on Jun 14, 2010

@author: oabalbin
'''
'''
Created on Oct 16, 2009

@author: oabalbin
'''
import sys
import operator
import glob
import os
#import pickle
import numpy as np
import scipy as sp
from scipy import stats
#from optparse import OptionParser
#from datetime import datetime
from collections import deque, defaultdict


from amarillo.common.classes import myArray, gene_copa
#import signatures.db.gene_annotations as ga


def open_files(folderpath,filename):
    '''
    Open the file with the expression for the kinases and other genes.
    '''
    myfiles=[]
    for infile in glob.glob( os.path.join(folderpath, filename) ):
        myfiles.append(infile)
    
    return myfiles

def list_of_names(inputfile):
        
        genlist=deque()
        for line in inputfile:        
            line = line.strip('\n')
            fields = line.split('\t')
            #print fields 
            # To skip headers, star reading samples in column 7 of the file
            if fields[0][0] == '#': 
                continue                             
            if fields[0]=='':
                continue
            
            genlist.append(fields[0])
        
        return genlist

def dict_of_names(inputfile):
        
        genlist=defaultdict()
        i=0
        for line in inputfile:        
            line = line.strip('\n')
            fields = line.split('\t')
            #print fields 
            # To skip headers, star reading samples in column 7 of the file
            if fields[0][0] == '#': 
                continue                             
            if fields[0]=='':
                continue
            
            genlist[fields[0]] = i
            i+=1 
        
        return genlist


def create_myArray(myfile, median_threshold, normal_samples):
    '''
    It creates a sample object and calculates a z-score for the expression level of each gene
    in the sample. Z-score (expvalues - genemedian)/genesd
    Returns a geneArray and also a samplelist with the expression for each sample
    The input matrix is already lo2 normalized    
    '''
    gene_dictionary, sample_dictionary, gene_interval_dict = defaultdict(), defaultdict(), defaultdict()
    celllines =  deque()
    normal_samples_list=[]
    
    first=True

    for line in myfile:
        if line is None:
            continue 
        if line.startswith('#'):
            continue
        # Handeling headers ....
       
        if line.startswith('gene'):
            celllines = line.strip('\n').split('\t')[2:]
            for j,sp in enumerate(celllines):
                sample_dictionary[j] = sp
                print sp,j
                if sp in normal_samples:
                    normal_samples_list.append(j)
                    
            k=0
            continue
        
        # Reading the file.
        fields = line.strip('\n').split('\t') 
        gene_dictionary[k] = fields[0]
        gene_interval_dict[k] = fields[1]
        
        exp=np.zeros(len(fields[2:]))
        expmasked = np.ma.make_mask_none((len(fields[2:]),))        
        exp = np.array(map(float,fields[2:]))
        
        # shift the values by one. In order to avoid the zero coverage values which will be inf when log is applied.    
        expvalues = np.array([exp]) + 1.0
        expvalues = np.ma.array(expvalues,mask=[expmasked])
        
        if first==True:
            geneRawValues = np.array(expvalues)
            first=False
        else:
            geneRawValues = np.concatenate((geneRawValues,expvalues),axis=0)
        
        k+=1
    
    # Pre-process the matrix data in order to filter out low expressing genes.
    # Calculate the median expression value for each gene.
    gene2keep = filter_low_expressing_genes(geneRawValues,median_threshold)
    geneRawValues = geneRawValues[gene2keep,:]
    gene_dictionary = rebuild_gene_dictionary(gene_dictionary,gene2keep)
    gene_interval_dict = rebuild_gene_dictionary(gene_interval_dict,gene2keep)
    
    # Create an GeneArray object
    thisGeneArray = myArray(geneRawValues, gene_dictionary, sample_dictionary, normal_samples_list, gene_interval_dict)

    return thisGeneArray


def filter_low_expressing_genes(geneRawValues,threshold):
    """
    It filters out genes with very low expression values. 
    genes which max expression value  < threshold for all samples are removed.          
    """
    print geneRawValues.shape[0]
    gene_indicators = np.arange(geneRawValues.shape[0])
    exp_cutoff = np.ma.median(geneRawValues, axis=1)
    #median_cutoff = sp.stats.scoreatpercentile(gene_medians,threshold) 
    #exp_cutoff = np.ma.max(geneRawValues, axis=1)
    gene2keep = gene_indicators[exp_cutoff > threshold]
    print len(gene2keep)
    return gene2keep

def rebuild_gene_dictionary(genedict,genes2keep):
    """
    """
    new_gene_dictionary=defaultdict()
    for i in range(len(genes2keep)):
        new_gene_dictionary[i] = genedict[genes2keep[i]]
    
    return new_gene_dictionary
    

def specificity_sensitivity(normal_samples, outlier_samples, all_samples):
    """
    It calculates the specificity and sentivity of the outlier calling
    """
    # Specificity
    len_outlier_samples, len_normal_samples, len_all_samples =  float(len(outlier_samples)), float(len(normal_samples)), float(len(all_samples))
    false_positives = len(normal_samples.intersection(set(outlier_samples)))
    true_negatives = len_normal_samples - false_positives
    specificity = true_negatives / (true_negatives + false_positives)
    
    #Sensitivity
    true_positives = len_outlier_samples - false_positives
    false_negatives =  len_all_samples - len_normal_samples
    sensitivity = true_positives / (true_positives + false_negatives)
    #Prevalence
    prevalence = len_outlier_samples / (len_all_samples - len_normal_samples)
    
    return specificity, sensitivity, prevalence


def copa_at_different_percentiles(thisGeneArray,percentils_list):
    """
    Run copa a different percentiles
    Return the following dictionaries gene_outliersample_dict, copa_rank_dict, copa_score_dict
    """
    gene_outliersample_dict, copa_rank_dict, copa_score_dict = defaultdict(list), defaultdict(list), defaultdict(list)
    
    for percentil in percentils_list:
        gene_outliersample_dict, copa_rank_dict, copa_score_dict = thisGeneArray.get_the_outliers(percentil, gene_outliersample_dict, copa_rank_dict, copa_score_dict)
    
    return gene_outliersample_dict, copa_rank_dict, copa_score_dict


def summarize_results_copa_percentiles(thisGeneArray, gene_outliersample_dict, copa_rank_dict, copa_score_dict, percentils_list, sdv_factor, method='mean'):
    """
    Returns a rank list of final outliers.
    """
    mean_gene_copa_score_list=[]
    indicator=np.array(range(len(percentils_list)))
    gene_copascore_values = np.empty(len(copa_rank_dict))
    
    for i in  range(len(copa_rank_dict)):
        if method == 'median':
            mean_copa_rank, mean_copa_score = np.median(copa_rank_dict[i]), np.median(copa_score_dict[i])
        elif method == 'mean':
            mean_copa_rank, mean_copa_score = np.mean(copa_rank_dict[i]), np.mean(copa_score_dict[i])
        #elif method == 'sum':
            
        genename, mean_copa_ind = thisGeneArray.gendict[i], indicator[copa_score_dict[i]>=mean_copa_score][0]
        genelocation = thisGeneArray.genInterval[i]
        outlier_samples = gene_outliersample_dict[i][mean_copa_ind]
        # quantifying how many normal samples the outlier hits
        normal_samples = set(thisGeneArray.normal_samples)
        all_samples = thisGeneArray.sampledict.keys()
        # Calculation of specificity, sensitivy, and prevalence
        specificity, sensitivity, prevalence = specificity_sensitivity(normal_samples, outlier_samples, all_samples)
        quality_outliers = [specificity, sensitivity, prevalence]
        
        # Make a temporal array with only the copa score for each gene
        gene_copascore_values[i] = mean_copa_score
        # weigth the mean_copa_score for the number of samples in which that gene have a value greater than mean_copa_score
        mean_copa_score = mean_copa_score*len(outlier_samples)
        # Filter outliers according to the level of variance in the normals
        sdt_normals = thisGeneArray.sdv_normals(i)
        sdt_tumors = thisGeneArray.sdv_tumors_mets(i)
        if sdt_normals > sdt_tumors*sdv_factor:
            continue
        
        # Make a list of outliers  
        mean_gene_copa_score_list.append( gene_copa(i,genename,mean_copa_rank,mean_copa_score, mean_copa_ind, outlier_samples, \
                                                    quality_outliers, genelocation) )
        
    #sorted list of outliers
    final_list_genes = sorted(mean_gene_copa_score_list, key=operator.attrgetter('mean_copa_score'), reverse=True)
    
    return final_list_genes, gene_copascore_values


def summarize_results_copa_percentiles_ETS(thisGeneArray, gene_outliersample_dict, copa_rank_dict, copa_score_dict, percentils_list, sdv_factor, method='mean'):
    """
    Returns a rank list of final outliers.
    """
    mean_gene_copa_score_list=[]
    indicator=np.array(range(len(percentils_list)))
    gene_copascore_values = np.empty(len(copa_rank_dict))
    
    for i in  range(len(copa_rank_dict)):
        if method == 'median':
            mean_copa_rank, mean_copa_score = np.median(copa_rank_dict[i]), np.median(copa_score_dict[i])
        elif method == 'mean':
            mean_copa_rank, mean_copa_score = np.mean(copa_rank_dict[i]), np.mean(copa_score_dict[i])
        #elif method == 'sum':
            
        genename, mean_copa_ind = thisGeneArray.gendict[i], indicator[copa_score_dict[i]>=mean_copa_score][0]
        genelocation = thisGeneArray.genInterval[i]
        outlier_samples = gene_outliersample_dict[i][mean_copa_ind]
        # quantifying how many normal samples the outlier hits
        normal_samples = set(thisGeneArray.normal_samples)
        all_samples = thisGeneArray.sampledict.keys()
        # Calculation of specificity, sensitivy, and prevalence
        specificity, sensitivity, prevalence = specificity_sensitivity(normal_samples, outlier_samples, all_samples)
        quality_outliers = [specificity, sensitivity, prevalence]
        
        # Make a temporal array with only the copa score for each gene
        gene_copascore_values[i] = mean_copa_score
        # weigth the mean_copa_score for the number of samples in which that gene have a value greater than mean_copa_score
        mean_copa_score = mean_copa_score*len(outlier_samples)
        # Filter outliers according to the level of variance in the normals
        sdt_normals = thisGeneArray.sdv_normals(i)
        sdt_tumors = thisGeneArray.sdv_tumors_mets(i)
        false_positives = list(set(outlier_samples).intersection(normal_samples))

        if (len(outlier_samples) <=2) or (sdt_normals > sdt_tumors*sdv_factor) or (len(false_positives) > 2):
            continue
        
        # Make a list of outliers  
        mean_gene_copa_score_list.append( gene_copa(i,genename,mean_copa_rank,mean_copa_score, mean_copa_ind, outlier_samples, \
                                                    quality_outliers, genelocation) )
        
    #sorted list of outliers
    final_list_genes = sorted(mean_gene_copa_score_list, key=operator.attrgetter('mean_copa_score'), reverse=True)
    
    return final_list_genes, gene_copascore_values



def identify_outliers_in_sample(thisGeneArray,final_list_genes, gene_outliersample_dict, gene_copascore_values, outliers_cutoff, not_report_isoforms=False):
        
    outliers_in_samples=np.zeros((thisGeneArray.expVal.shape[0],thisGeneArray.expVal.shape[1]))
    outliers_in_samples_copascore = np.zeros((thisGeneArray.expVal.shape[0],thisGeneArray.expVal.shape[1]))    
    list_outliers_report = deque()
    list_unique_outliers = defaultdict()
    for genename in thisGeneArray.gendict.values():
        list_unique_outliers[genename]=0

    
    i=0
    for gene in final_list_genes: 
        if i > outliers_cutoff:
            break
        
        thisoutlier = [ i, gene.indicator, gene.name, gene.location, gene.mean_copa_score, np.median(thisGeneArray.expVal[gene.indicator]), 
                       np.max(thisGeneArray.expVal[gene.indicator]), np.min(thisGeneArray.expVal[gene.indicator]), gene.quality[0], gene.quality[1],
                       gene.quality[2] ]
        
        # exclude isoforms        
        if not_report_isoforms:
            if list_unique_outliers[gene.name] == 0:
                list_outliers_report.append(thisoutlier)
                list_unique_outliers[gene.name] = 1
                i+=1
            else:
                continue
        else:
            list_outliers_report.append(thisoutlier)
            i+=1
            
        ### Most important thing  of this function    
        # Get the samples in which this gene is an outlier
        samples = gene_outliersample_dict[gene.indicator][gene.mean_copa_ind]
        outliers_in_samples[gene.indicator,samples] = 1
        # Matrix with copa score for the outliers in each sample
        outliers_in_samples_copascore[gene.indicator,samples] = gene_copascore_values[gene.indicator]
    
    return list_outliers_report, outliers_in_samples, outliers_in_samples_copascore 



def identify_outliers_in_sample_simple(thisGeneArray,final_list_genes, gene_outliersample_dict, outliers_cutoff, not_report_isoforms=False):
        
    list_outliers_report = deque()
    list_unique_outliers = defaultdict()
    for genename in thisGeneArray.gendict.values():
        list_unique_outliers[genename]=0

    
    i=0
    for gene in final_list_genes: 
        if i > outliers_cutoff:
            break
        
        ### Most important thing  of this function    
        # Get the samples in which this gene is an outlier
        samples = gene_outliersample_dict[gene.indicator][gene.mean_copa_ind]
        
        thisoutlier = [ i, gene.indicator, gene.name, gene.location, gene.mean_copa_score, np.median(thisGeneArray.expVal[gene.indicator]), 
                       np.max(thisGeneArray.expVal[gene.indicator]), np.min(thisGeneArray.expVal[gene.indicator]), gene.quality[0], gene.quality[1],
                       gene.quality[2], samples ]
        
        # exclude isoforms        
        if not_report_isoforms:
            if list_unique_outliers[gene.name] == 0:
                list_outliers_report.append(thisoutlier)
                list_unique_outliers[gene.name] = 1
                i+=1
            else:
                continue
        else:
            list_outliers_report.append(thisoutlier)
            i+=1
                
    return list_outliers_report 





def print_list_of_outliers(list_outliers_report, outfile_outliers_list):
    """
    Prints the list of outliers
    """       
    
    for thisoutlier in list_outliers_report:
        print ",".join(map(str,thisoutlier)).replace(',','\t')
        outfile_outliers_list.write(",".join(map(str,thisoutlier)).replace(',','\t')+'\n')
        

def print_matrix_of_outliers(list_of_samples, list_outliers_report, outliers_in_samples_copascore, outfile_outliers_matrix):
    """
    Prints the matrix of outliers
    """       
    outfile_outliers_matrix.write("geneName"+'\t'+",".join(list_of_samples).replace(',','\t')+'\n')
    
    for thisoutlier in list_outliers_report:
        eline = list(outliers_in_samples_copascore[thisoutlier[1],:])
        outfile_outliers_matrix.write(thisoutlier[2]+'\t'+",".join(map(str,eline)).replace(',','\t')+'\n') 


def print_outliers_by_samples(thisGeneArray, gene_copascore_values, outliers_in_samples, outfile_outliersinsample, nout2report):
    """
    Prints a text file with the name of the outliers for each sample.
    """
    
    gene_names_array = np.array(thisGeneArray.get_list_of_genes())
        
    for j, sample in thisGeneArray.sampledict.iteritems():
        # gets the raw copa score (not the weighted by num of samples). Rank using the copa score
        #gene_expression = gene_copascore_values[outliers_in_samples[:,j]==1]
        gene_expression = thisGeneArray.expVal_copacohort[outliers_in_samples[:,j]==1,j]
        # sort outliers genes according to expression
        sorted_outliers = np.ma.argsort(gene_expression)
        # Get gene names    
        oulier_genes_insample = gene_names_array[outliers_in_samples[:,j]==1]
        outliers_gene_expr = [[str(gene), str(exp)] for gene, exp in zip(oulier_genes_insample[sorted_outliers[::-1]], gene_expression[sorted_outliers[::-1]])]
        outliers_gene_expr = sum(outliers_gene_expr,[])
        print j, sample, outliers_gene_expr[:nout2report]
                
        outfile_outliersinsample.write(str(j)+'\t'+sample+'\t'+",".join(outliers_gene_expr[:nout2report]).replace(',','\t')+'\n')


def get_outliers_matrix_by_samples(thisGeneArray, gene_copascore_values, outliers_in_samples, outfile_outliersinsample, 
                                   nout2report, cohort_dictionary):
    """
    Prints a text file with the name of the outliers for each sample.
    """
    nout2report2 = nout2report*2
    gene_names_array = np.array(thisGeneArray.get_list_of_genes())

    for j, sample in thisGeneArray.sampledict.iteritems():
        # gets the raw copa score (not the weighted by num of samples). Rank using the copa score
        #gene_expression = gene_copascore_values[outliers_in_samples[:,j]==1]
        gene_expression = thisGeneArray.expVal_copacohort[outliers_in_samples[:,j]==1,j]
        # sort outliers genes according to expression
        sorted_outliers = np.ma.argsort(gene_expression)
        # Get gene names    
        oulier_genes_insample = gene_names_array[outliers_in_samples[:,j]==1]
        outliers_gene_expr = [[str(gene), str(exp)] for gene, exp in zip(oulier_genes_insample[sorted_outliers[::-1]],                                              
                                                                         gene_expression[sorted_outliers[::-1]])]
        outliers_gene_expr = sum(outliers_gene_expr,[])  
        missing_outliers = nout2report2 - len(outliers_gene_expr)
        
        if missing_outliers > 0:
            print missing_outliers, nout2report2, len(outliers_gene_expr)
            for i in range(missing_outliers):
                outliers_gene_expr.append('NAN')
                
        cohort_dictionary[sample].append(outliers_gene_expr[:nout2report2])
                        
    return cohort_dictionary


def get_outliers_matrix_by_samples_mod(thisGeneArray, gene_copascore_values, outliers_in_samples, outfile_outliersinsample, 
                                   nout2report):
    """
    Prints a text file with the name of the outliers for each sample.
    """
    nout2report2 = nout2report*2
    gene_names_array = np.array(thisGeneArray.get_list_of_genes())
    cohort_dictionary=defaultdict()
    outliers_sample_intersection={}

    for j, sample in thisGeneArray.sampledict.iteritems():
        # set blank list
        outliers_gene_expr=[]
        # gets the raw copa score (not the weighted by num of samples). Rank using the copa score
        #gene_expression = gene_copascore_values[outliers_in_samples[:,j]==1]
        gene_expression = thisGeneArray.expVal_copacohort[outliers_in_samples[:,j]==1,j]
        # sort outliers genes according to expression
        sorted_outliers = np.ma.argsort(gene_expression)
        # Get gene names    
        oulier_genes_insample = gene_names_array[outliers_in_samples[:,j]==1]
        for gene, exp in zip(oulier_genes_insample[sorted_outliers[::-1]],gene_expression[sorted_outliers[::-1]]):
            outliers_gene_expr.append([str(gene), str(exp)])
            print gene,exp
            try:
                myexp = outliers_sample_intersection[gene]
                if myexp < exp:
                    outliers_sample_intersection[gene] = exp
                
            except KeyError:
                outliers_sample_intersection[gene] = exp

        #outliers_gene_expr = sum(outliers_gene_expr,[])  
        missing_outliers = nout2report - len(outliers_gene_expr)
        
        if missing_outliers > 0:
            for i in range(missing_outliers):
                outliers_gene_expr.append(['NAN','NAN'])
                
        cohort_dictionary[sample]=outliers_gene_expr[:nout2report2]
    
    outliers_intersection_list = sorted(outliers_sample_intersection.iteritems(), key=operator.itemgetter(1), reverse=True)
    print outliers_intersection_list
    outliers_intersection_list_names={}
    for i in range(len(outliers_intersection_list)):
        gen = outliers_intersection_list[i][0]
        outliers_intersection_list_names[gen]=i
        if gen=='ERBB2':
            print i
    
    return cohort_dictionary, outliers_intersection_list_names



def print_copa_matrix(thisGeneArray, outfile_outliers_matrix):
    """
    """
    list_of_samples = thisGeneArray.get_list_of_samples()
    #print list_of_samples
    list_of_genes = thisGeneArray.get_list_of_genes()
    outfile_outliers_matrix.write("geneName"+'\t'+",".join(list_of_samples).replace(',','\t')+'\n')
    
    for i in range(thisGeneArray.expVal_copacohort.shape[0]):
        eline = list(thisGeneArray.expVal_copacohort[i,:])
        outfile_outliers_matrix.write(list_of_genes[i]+'\t'+",".join(map(str,eline)).replace(',','\t')+'\n') 

def print_foldchange_matrix(thisGeneArray, outfile_outliers_matrix):
    """
    """
    list_of_samples = thisGeneArray.get_list_of_samples()
    #print list_of_samples
    list_of_genes = thisGeneArray.get_list_of_genes()
    outfile_outliers_matrix.write("geneName"+'\t'+",".join(list_of_samples).replace(',','\t')+'\n')
    
    for i in range(thisGeneArray.expVal_foldchange.shape[0]):
        eline = list(thisGeneArray.expVal_foldchange[i,:])
        outfile_outliers_matrix.write(list_of_genes[i]+'\t'+",".join(map(str,eline)).replace(',','\t')+'\n') 


def print_foldchange_outliers_by_samples(thisGeneArray,outfile_outliersinsample, nout2report):
    """
    Prints a text file with the name of the outliers for each sample.
    """
    
    gene_names_array = np.array(thisGeneArray.get_list_of_genes())
 
    for j, sample in thisGeneArray.sampledict.iteritems():
        # gets the raw fold change
        gene_expression = thisGeneArray.expVal_foldchange[:,j]
        # sort outliers genes according to expression
        sorted_outliers = np.ma.argsort(gene_expression)
        # Get gene names    
        oulier_genes_insample = gene_names_array
        outliers_gene_expr = [[str(gene), str(exp)] for gene, exp in zip(oulier_genes_insample[sorted_outliers[::-1]], gene_expression[sorted_outliers[::-1]])]
        outliers_gene_expr = sum(outliers_gene_expr,[])
        print j, sample, outliers_gene_expr[:nout2report]
        
        outfile_outliersinsample.write(str(j)+'\t'+sample+'\t'+",".join(outliers_gene_expr[:nout2report]).replace(',','\t')+'\n')


############################


        










